#include "driver.hpp"

// global counters and switches
double average_iter_num;
double average_error_num;
double average_correct_blk_num;
double average_error_blk_num;
double average_undet_blk_num;
double average_det_blk_num;
double average_final_ber;
int trial_num;
double fliprate;

const static char* mat_dir = (char*) "matrices/H1057.244"; //H816"; //H4376"; //H816";  // the directory of the parity check matrix

// find the index of the first "1" in a vector
long index_of_first_one(vec_GF2& v){
  for(long i = 1; i <= N; i++){
    if (IsOne(v(i)))
      return i;
  }
  cout << "Bug: the vector contains all zeros, the matrix is not full rank" << endl;
  exit(EXIT_FAILURE);
}

// FIXME: this will bring the cache effects on board
// switch the i-th and j-th columns of matrix m
void switch_column(mat_GF2& m, long i, long j){
  GF2 temp;
  for(long k = 1; k <= m.NumRows(); k++){
    temp = m(k, i);
    m(k, i) = m(k, j);
    m(k, j) = temp;
  }
}

// do column switches if the matrix shape is bad
void adjust_columns(mat_GF2& H, vector< pair<long, long> >& perm){
  long i;
  long j;
  for(i = 1; i <= M; i++){
    j = index_of_first_one(H(i));
    if (i != j) {
      // start switching columns
#ifdef DEBUG
	cout << "switching columns of H:" << i << " and " << j << endl;
#endif
      switch_column(H, i, j);
      perm.push_back(pair<long, long>(i, j));
    }
  }
}

// FIXME: works only for a full rank matrix. This function oes inplace
// conversion to compute the reduced row echelon form of matrix H.
void to_rref(mat_GF2& H){
  long i, j;
  for(i = 2; i <= M; i++){
    for(j = 1; j <= i - 1; j++){      
      if (IsOne(H(j, i))) {
	H(j) = H(j) + H(i);
      }
    }
  }
}

// get the submatrix from matrix m.
const mat_GF2 submatrix(mat_GF2& m, long i1, long i2, long j1, long j2){
  mat_GF2 result = mat_GF2();
  result.SetDims(i2-i1+1, j2-j1+1);
  long a, b, p, q;
  for(a = i1, p = 1; a <= i2; a++, p++){
    for(b = j1, q = 1; b <= j2; b++, q++){
      result(p, q) = m(a, b);
    }
  }
  return result;
}


// read in the parity-check matrix and store the matrix in a 1-d
// array.
void read_matrix_H(mat_GF2& H){
  int b;
  H.SetDims(M, N);  // M * N matrix
  ifstream infile(mat_dir);
  if (!infile) {
    cout << "The following file can not be opened:" << mat_dir << endl;
    exit(EXIT_FAILURE);
  }
  for(long i = 1; i <= M; i++){
    for(long j = 1; j <= N; j++){
      infile >> b;
      H(i, j) = to_GF2(b);
    }
  }
}


// compute the generator matrix from the parity check matrix
void compute_G(mat_GF2& G, mat_GF2& H){
  G.SetDims(K, N);
  mat_GF2 Hcopy = mat_GF2(H);
  vector< pair<long, long> > perms;
  gauss(Hcopy);
  adjust_columns(Hcopy, perms);  
  to_rref(Hcopy);
  mat_GF2 transposedMat = transpose(submatrix(Hcopy, 1, M, N - K + 1, N));
  long i, j;
  for(i = 1; i <= K; i++){
    for(j = 1; j <= M; j++){
      G(i, j) = transposedMat(i, j);
    }
  }
  // assign identity diagonal matrix to G
  for(i = 1; i <= K; i++){
    G(i, M + i) = to_GF2(1);
  }
  // perform inverse column permutation
  for(i = perms.size() - 1; i >= 0; i--){
    switch_column(G, perms[i].first, perms[i].second);
  }
}

void simulate_decoding(const mat_GF2& H, const mat_GF2& G){
  vec_GF2 codeword;
  vec_GF2 correctedBits;
  vec_GF2 original;
  vec_GF2 noise;
  int iternum;
  long errnum;
  double* initialProbs = new double[N+1];
  initialProbs[0] = 0;
  for(long i = 1; i <= trial_num; i++){
    iternum = 0;
    generate_random_noisy_codeword(original, codeword, G);
    for(long j = 1; j <= N; j++){
      if (IsZero(codeword(j)))
	initialProbs[j] = fliprate;
      else
	initialProbs[j] = 1 - fliprate;
    }

#ifdef VERBOSE
      cout << "Start the " << i << "-th round of experiment" << endl;
#endif
     decodeSPA(correctedBits, H, iternum, initialProbs); 
     noise = correctedBits - original;
     average_iter_num += iternum;
     errnum = weight(noise);
     average_error_num += static_cast<double>(errnum);
     average_final_ber += static_cast<double>(errnum)/static_cast<double>(N);
  }

#ifdef DEBUG
  cout << "The original codeword:" << endl;
  cout << original << endl;
  cout << "Print out the corrected bits:" << endl;
  cout << correctedBits << endl;
#endif



  delete [] initialProbs;

}

// initialize the performance counters
void init_performance_counters() {
  average_iter_num = 0;
  average_error_num = 0;
  average_correct_blk_num = 0;
  average_error_blk_num = 0;
  average_undet_blk_num = 0;
  average_det_blk_num = 0;
  average_final_ber = 0;
}

void compute_average_performance() {
  // average the performance counter
  average_iter_num /= trial_num;
  average_error_num /= trial_num;
  average_final_ber /= trial_num;
}


// prints out the information of the experimentation
void print_info(){
  cout << "Exp info" << endl;
  cout << "-----------------------------------" << endl;
  cout << "N:" << N << endl;
  cout << "M:" << M << endl;
  cout << "K:" << K << endl;
  cout << "Rate:" << static_cast<double>(K)/static_cast<double>(N) << endl;
  cout << "BSC flip rate:" << fliprate << endl;
  cout << "Trials:" << trial_num << endl;
  cout << "Max itern:" << MAX_ITER << endl;
  cout << "-----------------------------------" << endl;
}


// prints out the performance results
void print_statistics() {
  cout << "Statistics:" << endl;
  cout << "-----------------------------------" << endl;
  cout << "Ave itern:" << average_iter_num << endl;
  cout << "Ave err num:" << average_error_num << endl;
  cout << "Ave BER:" << average_final_ber << endl;
  cout << "-----------------------------------" << endl;
}

// main entry for testing the LDPC decoding
int main(int argc, char *argv[]) {  
  if(argc != 3){
    cout << "incorrect number of parameters" << endl;
    exit(EXIT_FAILURE);
  }
  trial_num = atoi(argv[1]);
  fliprate = atof(argv[2]);
  mat_GF2 H = mat_GF2();
  mat_GF2 G = mat_GF2();
  read_matrix_H(H);
  compute_G(G, H);
  print_info();
#ifdef DEBUG
    verify_G(G, H);
#endif
    init_performance_counters();
  simulate_decoding(H, G);
  compute_average_performance();
  print_statistics();
  //write_counters_to_file();
  return 0;
}
